#include <algorithm>
#include <cfloat>
#include <vector>

#include "caffe/layers/pooling_layer.hpp"
#include "caffe/util/math_functions.hpp"

namespace caffe {

template <typename Ftype, typename Itype>
__global__ void MaxPoolForward(const int nthreads,
    const Ftype* bottom_data, const int num, const int channels,
    const int height, const int width, const int pooled_height,
    const int pooled_width, const int kernel_h, const int kernel_w,
    const int stride_h, const int stride_w, const int pad_h, const int pad_w,
    Ftype* top_data, Itype* mask, Ftype* top_mask) {
  CUDA_KERNEL_LOOP(index, nthreads) {
    const int pw = index % pooled_width;
    const int ph = (index / pooled_width) % pooled_height;
    const int c = (index / pooled_width / pooled_height) % channels;
    const int n = index / pooled_width / pooled_height / channels;
    int hstart = ph * stride_h - pad_h;
    int wstart = pw * stride_w - pad_w;
    const int hend = min(hstart + kernel_h, height);
    const int wend = min(wstart + kernel_w, width);
    hstart = max(hstart, 0);
    wstart = max(wstart, 0);
    float maxval = -static_cast<float>(max_dtype<Ftype>());
    int maxidx = -1;
    const Ftype* const bottom_slice =
        bottom_data + (n * channels + c) * height * width;
    for (int h = hstart; h < hend; ++h) {
      for (int w = wstart; w < wend; ++w) {
        int idx = h * width + w;
        if (static_cast<float>(bottom_slice[idx]) > maxval) {
          maxidx = idx;
          maxval = bottom_slice[maxidx];
        }
      }
    }
    top_data[index] = maxval;
    if (mask) {
      mask[index] = maxidx;
    } else {
      top_mask[index] = maxidx;
    }
  }
}

template <typename Ftype>
__global__ void AvePoolForward(const int nthreads,
    const Ftype* const bottom_data, const int num, const int channels,
    const int height, const int width, const int pooled_height,
    const int pooled_width, const int kernel_h, const int kernel_w,
    const int stride_h, const int stride_w, const int pad_h, const int pad_w,
    Ftype* const top_data) {
  CUDA_KERNEL_LOOP(index, nthreads) {
    const int pw = index % pooled_width;
    const int ph = (index / pooled_width) % pooled_height;
    const int c = (index / pooled_width / pooled_height) % channels;
    const int n = index / pooled_width / pooled_height / channels;
    int hstart = ph * stride_h - pad_h;
    int wstart = pw * stride_w - pad_w;
    int hend = min(hstart + kernel_h, height + pad_h);
    int wend = min(wstart + kernel_w, width + pad_w);
    const int pool_size = (hend - hstart) * (wend - wstart);
    hstart = max(hstart, 0);
    wstart = max(wstart, 0);
    hend = min(hend, height);
    wend = min(wend, width);
    Ftype aveval = 0.;
    const Ftype* const bottom_slice =
        bottom_data + (n * channels + c) * height * width;
    for (int h = hstart; h < hend; ++h) {
      for (int w = wstart; w < wend; ++w) {
        aveval += (bottom_slice[h * width + w]/pool_size);
      }
    }
    top_data[index] = aveval;
    // top_data[index] = aveval / pool_size;
  }
}

template <typename Ftype>
__global__ void StoPoolForwardTrain(const int nthreads,
    const Ftype* const bottom_data,
    const int num, const int channels, const int height,
    const int width, const int pooled_height, const int pooled_width,
    const int kernel_h, const int kernel_w, const int stride_h,
    const int stride_w, Ftype* const rand_idx, Ftype* const top_data) {
  CUDA_KERNEL_LOOP(index, nthreads) {
    const int pw = index % pooled_width;
    const int ph = (index / pooled_width) % pooled_height;
    const int c = (index / pooled_width / pooled_height) % channels;
    const int n = index / pooled_width / pooled_height / channels;
    const int hstart = ph * stride_h;
    const int hend = min(hstart + kernel_h, height);
    const int wstart = pw * stride_w;
    const int wend = min(wstart + kernel_w, width);
    float cumsum = 0.F;
    const Ftype* const bottom_slice =
        bottom_data + (n * channels + c) * height * width;
    // First pass: get sum
    for (int h = hstart; h < hend; ++h) {
      for (int w = wstart; w < wend; ++w) {
        cumsum += bottom_slice[h * width + w];
      }
    }
    const Ftype thres = rand_idx[index] * cumsum;
    // Second pass: get value, and set index.
    cumsum = 0;
    for (int h = hstart; h < hend; ++h) {
      for (int w = wstart; w < wend; ++w) {
        cumsum += bottom_slice[h * width + w];
        if (cumsum >= thres) {
          rand_idx[index] = ((n * channels + c) * height + h) * width + w;
          top_data[index] = bottom_slice[h * width + w];
          return;
        }
      }
    }
  }
}


template <typename Ftype>
__global__ void StoPoolForwardTest(const int nthreads,
    const Ftype* const bottom_data,
    const int num, const int channels, const int height,
    const int width, const int pooled_height, const int pooled_width,
    const int kernel_h, const int kernel_w, const int stride_h,
    const int stride_w, Ftype* const top_data) {
  CUDA_KERNEL_LOOP(index, nthreads) {
    const int pw = index % pooled_width;
    const int ph = (index / pooled_width) % pooled_height;
    const int c = (index / pooled_width / pooled_height) % channels;
    const int n = index / pooled_width / pooled_height / channels;
    const int hstart = ph * stride_h;
    const int hend = min(hstart + kernel_h, height);
    const int wstart = pw * stride_w;
    const int wend = min(wstart + kernel_w, width);
    // We set cumsum to be 0 to avoid divide-by-zero problems
    Ftype cumsum = min_dtype<Ftype>();
    Ftype cumvalues = 0.;
    const Ftype* const bottom_slice =
        bottom_data + (n * channels + c) * height * width;
    // First pass: get sum
    for (int h = hstart; h < hend; ++h) {
      for (int w = wstart; w < wend; ++w) {
        cumsum += bottom_slice[h * width + w];
        cumvalues += bottom_slice[h * width + w] * bottom_slice[h * width + w];
      }
    }
    top_data[index] = cumvalues / cumsum;
  }
}


template <typename Ftype, typename Btype>
void PoolingLayer<Ftype, Btype>::Forward_gpu(const vector<Blob*>& bottom,
      const vector<Blob*>& top) {
  const Ftype* bottom_data = bottom[0]->gpu_data<Ftype>();
  int count = top[0]->count();
  // We'll output the mask to top[1] if it's of size >1.
  const bool use_top_mask = top.size() > 1;
  int* mask = NULL;
  Ftype* top_mask = NULL;
  cudaStream_t stream = Caffe::thread_stream();
  switch (this->layer_param_.pooling_param().pool()) {
  case PoolingParameter_PoolMethod_MAX:
    if (use_top_mask) {
      top_mask = top[1]->mutable_gpu_data<Ftype>();
    } else {
      mask = max_idx_.mutable_gpu_data();
    }
    // NOLINT_NEXT_LINE(whitespace/operators)
    MaxPoolForward<Ftype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
        count, bottom_data, bottom[0]->num(), channels_,
        height_, width_, pooled_height_, pooled_width_, kernel_h_,
        kernel_w_, stride_h_, stride_w_, pad_h_, pad_w_, top[0]->mutable_gpu_data<Ftype>(),
        mask, top_mask);
    break;
  case PoolingParameter_PoolMethod_AVE:
    // NOLINT_NEXT_LINE(whitespace/operators)
    AvePoolForward<Ftype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
        count, bottom_data, bottom[0]->num(), channels_,
        height_, width_, pooled_height_, pooled_width_, kernel_h_,
        kernel_w_, stride_h_, stride_w_, pad_h_, pad_w_, top[0]->mutable_gpu_data<Ftype>());
    break;
  case PoolingParameter_PoolMethod_STOCHASTIC:
    if (this->phase_ == TRAIN) {
      // We need to create the random index as well.
      caffe_gpu_rng_uniform(count, Ftype(0), Ftype(1), rand_idx_->mutable_gpu_data<Ftype>());
      // NOLINT_NEXT_LINE(whitespace/operators)
      StoPoolForwardTrain<Ftype><<<CAFFE_GET_BLOCKS(count),
                                   CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
          count, bottom_data, bottom[0]->num(), channels_,
          height_, width_, pooled_height_, pooled_width_, kernel_h_,
          kernel_w_, stride_h_, stride_w_,
          rand_idx_->mutable_gpu_data<Ftype>(), top[0]->mutable_gpu_data<Ftype>());
    } else {
      // NOLINT_NEXT_LINE(whitespace/operators)
      StoPoolForwardTest<Ftype><<<CAFFE_GET_BLOCKS(count),
                                  CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
          count, bottom_data, bottom[0]->num(), channels_,
          height_, width_, pooled_height_, pooled_width_, kernel_h_,
          kernel_w_, stride_h_, stride_w_, top[0]->mutable_gpu_data<Ftype>());
    }
    break;
  default:
    LOG(FATAL) << "Unknown pooling method.";
  }
  CUDA_POST_KERNEL_CHECK;
  CUDA_CHECK(cudaStreamSynchronize(stream));
}

template <typename Btype, typename Itype>
__global__ void MaxPoolBackward(const int nthreads, const Btype* top_diff,
    const Itype* mask, const Btype* top_mask, const int num,
    const int channels, const int height, const int width,
    const int pooled_height, const int pooled_width, const int kernel_h,
    const int kernel_w, const int stride_h, const int stride_w, const int pad_h,
    const int pad_w, Btype* bottom_diff) {
  CUDA_KERNEL_LOOP(index, nthreads) {
    // find out the local index
    // find out the local offset
    const int w = index % width;
    const int h = (index / width) % height;
    const int c = (index / width / height) % channels;
    const int n = index / width / height / channels;
    const int phstart =
         (h + pad_h < kernel_h) ? 0 : (h + pad_h - kernel_h) / stride_h + 1;
    const int phend = min((h + pad_h) / stride_h + 1, pooled_height);
    const int pwstart =
         (w + pad_w < kernel_w) ? 0 : (w + pad_w - kernel_w) / stride_w + 1;
    const int pwend = min((w + pad_w) / stride_w + 1, pooled_width);
    float gradient = 0;
    const int offset = (n * channels + c) * pooled_height * pooled_width;
    const Btype* const top_diff_slice = top_diff + offset;
    if (mask) {
      const int* const mask_slice = mask + offset;
      for (int ph = phstart; ph < phend; ++ph) {
        for (int pw = pwstart; pw < pwend; ++pw) {
          int idx = ph * pooled_width + pw;
          if (mask_slice[idx] == h * width + w) {
            gradient += static_cast<float>(top_diff_slice[idx]);
          }
        }
      }
    } else {
      const Btype* const top_mask_slice = top_mask + offset;
      for (int ph = phstart; ph < phend; ++ph) {
        for (int pw = pwstart; pw < pwend; ++pw) {
          int idx = ph * pooled_width + pw;
          if (static_cast<int>(top_mask_slice[idx]) == h * width + w) {
            gradient += static_cast<float>(top_diff_slice[idx]);
          }
        }
      }
    }
    bottom_diff[index] = gradient;
  }
}

template <typename Btype>
__global__ void AvePoolBackward(const int nthreads, const Btype* const top_diff,
    const int num, const int channels, const int height,
    const int width, const int pooled_height, const int pooled_width,
    const int kernel_h, const int kernel_w, const int stride_h,
    const int stride_w, const int pad_h, const int pad_w,
    Btype* const bottom_diff) {
  CUDA_KERNEL_LOOP(index, nthreads) {
    // find out the local index
    // find out the local offset
    const int w = index % width + pad_w;
    const int h = (index / width) % height + pad_h;
    const int c = (index / width / height) % channels;
    const int n = index / width / height / channels;
    const int phstart = (h < kernel_h) ? 0 : (h - kernel_h) / stride_h + 1;
    const int phend = min(h / stride_h + 1, pooled_height);
    const int pwstart = (w < kernel_w) ? 0 : (w - kernel_w) / stride_w + 1;
    const int pwend = min(w / stride_w + 1, pooled_width);
//    Btype gradient = 0;
    float gradient = 0.;
    const Btype* const top_diff_slice =
        top_diff + (n * channels + c) * pooled_height * pooled_width;
    for (int ph = phstart; ph < phend; ++ph) {
      for (int pw = pwstart; pw < pwend; ++pw) {
        // figure out the pooling size
        int hstart = ph * stride_h - pad_h;
        int wstart = pw * stride_w - pad_w;
        int hend = min(hstart + kernel_h, height + pad_h);
        int wend = min(wstart + kernel_w, width + pad_w);
        int pool_size = (hend - hstart) * (wend - wstart);
        gradient += static_cast<float>(top_diff_slice[ph * pooled_width + pw]) / pool_size;
      }
    }
    bottom_diff[index] = gradient;
  }
}

template <typename Btype>
__global__ void StoPoolBackward(const int nthreads,
    const Btype* const rand_idx, const Btype* const top_diff,
    const int num, const int channels, const int height,
    const int width, const int pooled_height, const int pooled_width,
    const int kernel_h, const int kernel_w, const int stride_h,
    const int stride_w, Btype* const bottom_diff) {
  CUDA_KERNEL_LOOP(index, nthreads) {
    // find out the local index
    // find out the local offset
    const int w = index % width;
    const int h = (index / width) % height;
    const int c = (index / width / height) % channels;
    const int n = index / width / height / channels;
    const int phstart = (h < kernel_h) ? 0 : (h - kernel_h) / stride_h + 1;
    const int phend = min(h / stride_h + 1, pooled_height);
    const int pwstart = (w < kernel_w) ? 0 : (w - kernel_w) / stride_w + 1;
    const int pwend = min(w / stride_w + 1, pooled_width);
    float gradient = 0.F;
    const Btype* const rand_idx_slice =
        rand_idx + (n * channels + c) * pooled_height * pooled_width;
    const Btype* const top_diff_slice =
        top_diff + (n * channels + c) * pooled_height * pooled_width;
    for (int ph = phstart; ph < phend; ++ph) {
      for (int pw = pwstart; pw < pwend; ++pw) {
        gradient += top_diff_slice[ph * pooled_width + pw] *
            (index == static_cast<int>(rand_idx_slice[ph * pooled_width + pw]));
      }
    }
    bottom_diff[index] = gradient;
  }
}


template <typename Ftype, typename Btype>
void PoolingLayer<Ftype, Btype>::Backward_gpu(const vector<Blob*>& top,
      const vector<bool>& propagate_down, const vector<Blob*>& bottom) {
  if (!propagate_down[0]) {
    return;
  }
  const Btype* top_diff = top[0]->gpu_diff<Btype>();
  Btype* bottom_diff = bottom[0]->mutable_gpu_diff<Btype>();
  const Btype* top_data = top[0]->gpu_data<Btype>();
  const Btype* bottom_data = bottom[0]->gpu_data<Btype>();
  const int count = bottom[0]->count();
  caffe_gpu_set<Btype>(count, Btype(0.), bottom_diff);
  // We'll output the mask to top[1] if it's of size >1.
  const bool use_top_mask = top.size() > 1;
  const int* mask = NULL;
  const Btype* top_mask = NULL;
  cudaStream_t stream = Caffe::thread_stream();
  switch (this->layer_param_.pooling_param().pool()) {
  case PoolingParameter_PoolMethod_MAX:
    if (use_top_mask) {
      top_mask = top[1]->gpu_data<Btype>();
    } else {
      mask = max_idx_.gpu_data();
    }
    {
      // NOLINT_NEXT_LINE(whitespace/operators)
      MaxPoolBackward<Btype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
          count, top_diff, mask, top_mask, top[0]->num(), channels_,
          height_, width_, pooled_height_, pooled_width_,
          kernel_h_, kernel_w_, stride_h_, stride_w_, pad_h_, pad_w_,
          bottom_diff);
    }
    break;
  case PoolingParameter_PoolMethod_AVE:
    // NOLINT_NEXT_LINE(whitespace/operators)
    AvePoolBackward<Btype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
        count, top_diff, top[0]->num(), channels_,
        height_, width_, pooled_height_, pooled_width_, kernel_h_,
        kernel_w_, stride_h_, stride_w_, pad_h_, pad_w_, bottom_diff);
    break;
  case PoolingParameter_PoolMethod_STOCHASTIC:
    // NOLINT_NEXT_LINE(whitespace/operators)
    StoPoolBackward<Btype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
        count, rand_idx_->gpu_data<Btype>(), top_diff,
        top[0]->num(), channels_, height_, width_, pooled_height_,
        pooled_width_, kernel_h_, kernel_w_, stride_h_, stride_w_,
        bottom_diff);
    break;
  default:
    LOG(FATAL) << "Unknown pooling method.";
  }
  CUDA_POST_KERNEL_CHECK;
  CUDA_CHECK(cudaStreamSynchronize(stream));
}

  // Auxiliary boundary-based Forward_gpu/Backward_gpu functions.
  // Added by Hao FU.
  template <typename Ftype, typename Btype>
  void PoolingLayer<Ftype, Btype>::Forward_gpu(const vector<Blob*>& bottom,
                                               const vector<Blob*>& top,
                                               BatchBoundary_t& boundary) {
    const int top_stride_count = top[0]->count(1);
    CHECK_EQ(top_stride_count, channels_ * pooled_height_ * pooled_width_)
      << "top_stride_count should be equal to channels_ * pooled_height_ * pooled_width_";
    const int bottom_stride_count = bottom[0]->count(1);
    int count = top_stride_count * (boundary.end_idx_ - boundary.start_idx_ + 1);
    const Ftype* bottom_data = bottom[0]->gpu_data<Ftype>() +
      (boundary.start_idx_ * bottom_stride_count);
    // We'll output the mask to top[1] if it's of size > 1.
    const bool use_top_mask = top.size() > 1;
    int* mask = nullptr;
    Ftype* top_mask = nullptr;
    cudaStream_t stream = Caffe::thread_stream();
    switch (this->layer_param_.pooling_param().pool()) {
    case PoolingParameter_PoolMethod_MAX:
      if (use_top_mask) {
        top_mask = top[1]->mutable_gpu_data<Ftype>() +
          (boundary.start_idx_ * top_stride_count);
      } else {
        mask = max_idx_.mutable_gpu_data() +
          (boundary.start_idx_ * top_stride_count);
      }
      // NOLINT_NEXT_LINE(whitespace/operators)
      MaxPoolForward<Ftype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
          count, bottom_data, (boundary.end_idx_ - boundary.start_idx_ + 1), channels_,
          height_, width_, pooled_height_, pooled_width_, kernel_h_,
          kernel_w_, stride_h_, stride_w_, pad_h_, pad_w_, top[0]->mutable_gpu_data<Ftype>() + (boundary.start_idx_ * top_stride_count),
          mask, top_mask);
      break;
    case PoolingParameter_PoolMethod_AVE:
      AvePoolForward<Ftype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
          count, bottom_data, (boundary.end_idx_ - boundary.start_idx_ + 1), channels_,
          height_, width_, pooled_height_, pooled_width_, kernel_h_,
          kernel_w_, stride_h_, stride_w_, pad_h_, pad_w_, top[0]->mutable_gpu_data<Ftype>() + (boundary.start_idx_ * top_stride_count));
      break;
    case PoolingParameter_PoolMethod_STOCHASTIC:
      if (this->phase_ == TRAIN) {
        // We need to create the random index as well.
        caffe_gpu_rng_uniform(count, Ftype(0), Ftype(1), rand_idx_->mutable_gpu_data<Ftype>() + (boundary.start_idx_ * top_stride_count));
        // NOLINT_NEXT_LINE(whitespace/operators)
        StoPoolForwardTrain<Ftype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
            count, bottom_data, (boundary.end_idx_ - boundary.start_idx_ + 1), channels_,
            height_, width_, pooled_height_, pooled_width_,
            kernel_h_, kernel_w_, stride_h_, stride_w_,
            rand_idx_->mutable_gpu_data<Ftype>() + (boundary.start_idx_ * top_stride_count),
            top[0]->mutable_gpu_data<Ftype>() + (boundary.start_idx_ * top_stride_count));
      } else {
        // NOLINT_NEXT_LINE(whitespace/operators)
        StoPoolForwardTest<Ftype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
                       count, bottom_data, (boundary.end_idx_ - boundary.start_idx_ + 1), channels_,
                       height_, width_, pooled_height_, pooled_width_,
                       kernel_h_, kernel_w_, stride_h_, stride_w_,
                       top[0]->mutable_gpu_data<Ftype>() + (boundary.start_idx_ * top_stride_count));
      }
      break;
    default:
      LOG(FATAL) << "Unknown pooling method.";
    }
    CUDA_POST_KERNEL_CHECK;
    CUDA_CHECK(cudaStreamSynchronize(stream));
  }
  template <typename Ftype, typename Btype>
  void PoolingLayer<Ftype, Btype>::Backward_gpu(const vector<Blob*>& top,
                                                const vector<bool>& propagate_down,
                                                const vector<Blob*>& bottom,
                                                BatchBoundary_t& boundary) {
    if (!propagate_down[0]) {
      return ;
    }
    const int top_stride_count = top[0]->count(1);
    const int bottom_stride_count = bottom[0]->count(1);
    CHECK_EQ(top_stride_count, channels_ * pooled_height_ * pooled_width_)
      << "top_stride_count should be equal to " << channels_ << "*" << pooled_height_
      << "*" << pooled_width_;

    const Btype* top_diff = top[0]->gpu_diff<Btype>() + (boundary.start_idx_ * top_stride_count);
    const Btype* top_data = top[0]->gpu_data<Btype>() + (boundary.start_idx_ * top_stride_count);
    const Btype* bottom_data = bottom[0]->gpu_data<Btype>() +
      (boundary.start_idx_ * bottom_stride_count);
    Btype* bottom_diff = bottom[0]->mutable_gpu_diff<Btype>() +
      (boundary.start_idx_ * bottom_stride_count);

    const int count = bottom_stride_count * (boundary.end_idx_ - boundary.start_idx_ + 1);
    caffe_gpu_set<Btype>(count, Btype(0.), bottom_diff);

    // We'll output the mask to top[1] if it's of size > 1.
    const bool use_top_mask = top.size() > 1;
    const int* mask = nullptr;
    const Btype* top_mask = nullptr;
    cudaStream_t stream = Caffe::thread_stream();

    switch (this->layer_param_.pooling_param().pool()) {
    case PoolingParameter_PoolMethod_MAX:
      if (use_top_mask) {
        top_mask = top[1]->gpu_data<Btype>() + (boundary.start_idx_ * top_stride_count);
      } else {
        mask = max_idx_.gpu_data() + (boundary.start_idx_ * top_stride_count);
      }
      // NOLINT_NEXT_LINE(whitespace/operators)
      MaxPoolBackward<Btype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
          count, top_diff, mask, top_mask, (boundary.end_idx_ - boundary.start_idx_ + 1), channels_,
          height_, width_, pooled_height_, pooled_width_,
          kernel_h_, kernel_w_, stride_h_, stride_w_, pad_h_, pad_w_,
          bottom_diff);
      break;
    case PoolingParameter_PoolMethod_AVE:
      // NOLINT_NEXT_LINE(whitespace/operators)
      AvePoolBackward<Btype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
          count, top_diff, (boundary.end_idx_ - boundary.start_idx_ + 1), channels_,
          height_, width_, pooled_height_, pooled_width_, kernel_h_,
          kernel_w_, stride_h_, stride_w_, pad_h_, pad_w_, bottom_diff);
      break;
    case PoolingParameter_PoolMethod_STOCHASTIC:
      // NOLINT_NEXT_LINE(whitespace/operators)
      StoPoolBackward<Btype><<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS, 0, stream>>>(
          count, rand_idx_->gpu_data<Btype>() + (boundary.start_idx_ * top_stride_count), top_diff,
          (boundary.end_idx_ - boundary.start_idx_ + 1), channels_, height_, width_, pooled_height_,
          pooled_width_, kernel_h_, kernel_w_, stride_h_, stride_w_,
          bottom_diff);
      break;
    default:
      LOG(FATAL)  << "Unknown pooling method.";
    }

    CUDA_POST_KERNEL_CHECK;
    CUDA_CHECK(cudaStreamSynchronize(stream));
  }
  // End. Added by Hao FU.

INSTANTIATE_LAYER_GPU_FUNCS_FB(PoolingLayer);

// Auxiliary Forward_gpu/Backward_gpu functions for the HGP4NN framework.
// Added by Hao FU.
INSTANTIATE_LAYER_GPU_BOUNDARY_FUNCS_FB(PoolingLayer);
// End. Added by Hao FU.

}  // namespace caffe
